#!/usr/bin/env python

import os
import yaml
import rospy
import tf2_ros
from std_srvs.srv import Trigger, TriggerResponse
from geometry_msgs.msg import TransformStamped

### @brief Class to manage the publishing of static transforms
### @details - In ROS, there can only be one static transform publisher per node
###            so any Python module that wants to publish a static transform publishes
###            the transform to this node via a topic; then this node aggregates and sends all the transforms
class StaticTransformManager(object):
    def __init__(self):
        self.transform_list = []
        self.br = tf2_ros.StaticTransformBroadcaster()
        
        # get save/load params
        load_transforms = rospy.get_param(
            "~load_transforms", 
            default=True)
        save_transforms = rospy.get_param(
            "~save_transforms", 
            default=True)

        # if save or load, do file and directory management
        if save_transforms or load_transforms:
            filename = rospy.get_param(
                "~transform_filename",
                default="static_transforms.yaml")
            self.filepath = rospy.get_param(
                "~transform_filepath", 
                default=os.path.join("/interbotix/static_transforms/", filename))
            filedir = os.path.dirname(self.filepath)
            self.full_filepath = os.path.join(os.getcwd(), self.filepath)
            rospy.loginfo("Transform filepath set to " + self.full_filepath)

            try: # check if directory exists, if not make one
                if not os.path.exists(filedir):
                    os.makedirs(filedir)
                    rospy.loginfo("Creating directory: " + os.path.join(os.getcwd(), filedir))
            except Exception: # if we fail (permissions, etc.)
                rospy.logwarn("Failed to create directory: " + e + " Transform file will not be saved.")

            # load if we need to
            if load_transforms:
                self.load_transforms()
        
            # save and set up saving service if we need to
            if save_transforms:
                rospy.Service('save_transforms', Trigger, self.trigger_cb)
                rospy.on_shutdown(self.save_transforms)

        rospy.Subscriber("static_transforms", TransformStamped, self.transform_cb)

        print("Initialized Static Transform Publisher!")

    ### @brief ROS Subscriber Callback that receives TransformStamped messages
    ### @param msg - TransformStamped ROS message
    ### @details - any new static transform is appended to the self.transform_list variable;
    ###            if the transform already exists, the old one is replaced with the new one
    def transform_cb(self, msg):
        rospy.loginfo(
            "Static Transform Publisher received TF from \"" + msg.header.frame_id + "\" to \"" + msg.child_frame_id + "\"")
        for indx in range(len(self.transform_list)):
            if msg.child_frame_id == self.transform_list[indx].child_frame_id:
                self.transform_list.pop(indx)
                break
        self.transform_list.append(msg)
        self.br.sendTransform(self.transform_list)

    ### @brief Load the transforms from the file (located at 'self.transform_filepath') and send them
    def load_transforms(self):
        try:
            with open(self.filepath, "r") as yamlfile:
                trans_list = yaml.safe_load(yamlfile)
            rospy.loginfo("Loaded static transforms from " + self.full_filepath)

            self.transform_list = []
            for trans_dict in trans_list:
                trans = TransformStamped()
                trans.header.frame_id = trans_dict["frame_id"]
                trans.child_frame_id = trans_dict["child_frame_id"]
                trans.transform.translation.x = trans_dict["x"]
                trans.transform.translation.y = trans_dict["y"]
                trans.transform.translation.z = trans_dict["z"]
                trans.transform.rotation.x = trans_dict["qx"]
                trans.transform.rotation.y = trans_dict["qy"]
                trans.transform.rotation.z = trans_dict["qz"]
                trans.transform.rotation.w = trans_dict["qw"]
                trans.header.stamp = rospy.Time.now()
                self.transform_list.append(trans)
            self.br.sendTransform(self.transform_list)

        except IOError:
            rospy.logwarn("File at %s does not exist yet. No static transforms loaded." % self.full_filepath)

    ### @brief At node shutdown, save all static transforms to the YAML file specified at 'self.transform_filepath'
    ### @details - if no transforms were published to this node, no file is created and no transforms are saved
    def save_transforms(self):
        if self.transform_list == []:
            return
        
        trans_list = []
        for trans in self.transform_list:
            trans_dict = {}
            trans_dict["frame_id"] = trans.header.frame_id
            trans_dict["child_frame_id"] = trans.child_frame_id
            trans_dict["x"] = float(trans.transform.translation.x)
            trans_dict["y"] = float(trans.transform.translation.y)
            trans_dict["z"] = float(trans.transform.translation.z)
            trans_dict["qx"] = float(trans.transform.rotation.x)
            trans_dict["qy"] = float(trans.transform.rotation.y)
            trans_dict["qz"] = float(trans.transform.rotation.z)
            trans_dict["qw"] = float(trans.transform.rotation.w)
            trans_list.append(trans_dict)

        with open(self.filepath, "w") as yamlfile:
            yaml.dump(trans_list, yamlfile, default_flow_style=False)

        rospy.loginfo("Saved static transforms to: " + self.full_filepath)

    ### @brief Service callback to save transforms now, instead of on node shutdown
    ### @returns - TriggerResponse stating that the TFs were saved
    def trigger_cb(self, req):
        try:
            self.save_transforms()
            return TriggerResponse(
                success=True,
                message="Static transforms saved successfully.")
        except Exception:
            rospy.logerr("Something went wrong when saving static transforms.")
            return TriggerResponse(
                success=False,
                message="Static transforms NOT saved successfully.")


def main():
    rospy.init_node('static_trans_pub')
    StaticTransformManager()
    rospy.spin()

if __name__=='__main__':
    main()
